package top.yudoge.dogeeditor.llm.util

import java.math.BigDecimal
import java.math.RoundingMode

object VectorUtils {

    fun cosSimilarity(vec1: List<Number>, vec2: List<Number>): BigDecimal {
        // 计算两个向量夹角余弦
        val vec1Magnitude = magnitude(vec1)
        val vec2Magnitude = magnitude(vec2)
        val dotProduct = BigDecimal.valueOf(vec1.zip(vec2).sumOf { (a, b) -> a.toDouble() * b.toDouble() })
        val magnitudeProduct = vec1Magnitude.multiply(vec2Magnitude)
        if (magnitudeProduct.compareTo(BigDecimal.ZERO) == 0) {
            return BigDecimal.ZERO // 无相似度
        }
        return dotProduct.divide(magnitudeProduct, 10, RoundingMode.HALF_UP)
    }

    fun bestSimilar(target: List<Number>, matches: List<List<Number>>): Int {
        if (matches.isEmpty()) {
            return -1
        }

        return matches.mapIndexed { index, match ->
            cosSimilarity(target, match) to index
        }.maxByOrNull { it.first }?.second!!
    }

    fun topKSimilar(target: List<Number>, matches: List<List<Number>>, k: Int): List<Int> {
        if (matches.isEmpty()) {
            return emptyList()
        }
        return matches.mapIndexed { index, match ->
            cosSimilarity(target, match) to index
        }.sortedByDescending { it.first }.take(k).map { it.second }.toList()
    }

    private fun magnitude(vec: List<Number>): BigDecimal {
        var powSum: Double = 0.0
        for (i in vec.indices) {
            powSum += Math.pow(vec[i].toDouble(), 2.0)
        }
        return BigDecimal.valueOf(Math.sqrt(powSum))
    }



}